// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/connectors/hive/hive_partition_util.h>

namespace kumo::pollux::connector::hive {

#define PARTITION_TYPE_DISPATCH(TEMPLATE_FUNC, typeKind, ...)               \
  [&]() {                                                                   \
    switch (typeKind) {                                                     \
      case TypeKind::BOOLEAN:                                               \
      case TypeKind::TINYINT:                                               \
      case TypeKind::SMALLINT:                                              \
      case TypeKind::INTEGER:                                               \
      case TypeKind::BIGINT:                                                \
      case TypeKind::VARCHAR:                                               \
      case TypeKind::VARBINARY:                                             \
        return POLLUX_DYNAMIC_SCALAR_TYPE_DISPATCH(                          \
            TEMPLATE_FUNC, typeKind, __VA_ARGS__);                          \
      default:                                                              \
        POLLUX_UNSUPPORTED(                                                  \
            "Unsupported partition type: {}", mapTypeKindToName(typeKind)); \
    }                                                                       \
  }()

namespace {
template <typename T>
inline std::string makePartitionValueString(T value) {
  return melon::to<std::string>(value);
}

template <>
inline std::string makePartitionValueString(bool value) {
  return value ? "true" : "false";
}

template <TypeKind Kind>
std::pair<std::string, std::string> makePartitionKeyValueString(
    const BaseVector* partitionVector,
    vector_size_t row,
    const std::string& name,
    bool isDate) {
  using T = typename TypeTraits<Kind>::NativeType;
  if (partitionVector->as<SimpleVector<T>>()->is_null_at(row)) {
    return std::make_pair(name, "");
  }
  if (isDate) {
    return std::make_pair(
        name,
        DATE()->toString(
            partitionVector->as<SimpleVector<int32_t>>()->value_at(row)));
  }
  return std::make_pair(
      name,
      makePartitionValueString(
          partitionVector->as<SimpleVector<T>>()->value_at(row)));
}

} // namespace

std::vector<std::pair<std::string, std::string>> extractPartitionKeyValues(
    const RowVectorPtr& partitionsVector,
    vector_size_t row) {
  std::vector<std::pair<std::string, std::string>> partitionKeyValues;
  for (auto i = 0; i < partitionsVector->childrenSize(); i++) {
    partitionKeyValues.push_back(PARTITION_TYPE_DISPATCH(
        makePartitionKeyValueString,
        partitionsVector->childAt(i)->type_kind(),
        partitionsVector->childAt(i)->loaded_vector(),
        row,
        as_row_type(partitionsVector->type())->nameOf(i),
        partitionsVector->childAt(i)->type()->isDate()));
  }
  return partitionKeyValues;
}

} // namespace kumo::pollux::connector::hive
